from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.callbacks import LlamaDebugHandler, CallbackManager
from llama_index.core.indices.common.struct_store.sql import SQLStructDatapointExtractor
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import  TextNode
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer, \
    DocumentSummaryIndex, SimpleDirectoryReader, VectorStoreIndex
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel
from llama_index.core.indices.property_graph.base import PropertyGraphIndex
from llama_index.core.indices.property_graph.retriever import PGRetriever
from llama_index.core.indices.property_graph.sub_retrievers.base import BasePGRetriever
from llama_index.core.indices.property_graph.sub_retrievers.custom import (
    CustomPGRetriever,
    CUSTOM_RETRIEVE_TYPE,
)
from llama_index.core.indices.property_graph.sub_retrievers.cypher_template import (
    CypherTemplateRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.llm_synonym import (
    LLMSynonymRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.text_to_cypher import (
    TextToCypherRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.vector import (
    VectorContextRetriever,
)
from llama_index.core.indices.property_graph.transformations.implicit import (
    ImplicitPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.schema_llm import (
    SchemaLLMPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.simple_llm import (
    SimpleLLMPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.dynamic_llm import (
    DynamicLLMPathExtractor,
)
from llama_index.core.indices.property_graph.utils import default_parse_triplets_fn

embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm

from llama_index.core.response_synthesizers.accumulate import Accumulate
from llama_index.core.response_synthesizers.base import BaseSynthesizer
from llama_index.core.response_synthesizers.compact_and_refine import (
    CompactAndRefine,
)
from llama_index.core.response_synthesizers.factory import get_response_synthesizer
from llama_index.core.response_synthesizers.generation import Generation
from llama_index.core.response_synthesizers.refine import Refine
from llama_index.core.response_synthesizers.simple_summarize import SimpleSummarize
from llama_index.core.response_synthesizers.tree_summarize import TreeSummarize
from llama_index.core.response_synthesizers.type import ResponseMode


llama_debug = LlamaDebugHandler(print_trace_on_end=True)
callback_manager = CallbackManager([llama_debug])

Settings.callback_manager = callback_manager

'''
simpleSummarize =SimpleSummarize(callback_manager=callback_manager)

rs=simpleSummarize.get_response(query_str="太阳的体积比地球的体积大多少",text_chunks=["太阳的体积10000立方米","地球的体积800立方米"])
print(rs)

'''

def acc():
    simpleSummarize =Accumulate(callback_manager=callback_manager)

    rs=simpleSummarize.get_response(query_str="Refine模式工作原理",text_chunks=['''Refine模式工作原理
    ‌初始回答生成‌：使用第一个检索到的文本块和text_qa_template生成初始回答
    ''','''Refine模式工作原理
    ‌‌迭代优化‌：将初始回答与后续文本块结合，通过refine_template提示词逐步优化答案，直至处理完所有相关文本块
    ''','''Refine模式工作原理
    上下文溢出处理‌：若单个文本块过长，会自动分割并继续迭代优化'''])

    print(rs)

def acc():
    simpleSummarize =Accumulate(callback_manager=callback_manager)

    rs=simpleSummarize.get_response(query_str="Refine模式工作原理",text_chunks=['''Refine模式工作原理
    ‌初始回答生成‌：使用第一个检索到的文本块和text_qa_template生成初始回答
    ''','''Refine模式工作原理
    ‌‌迭代优化‌：将初始回答与后续文本块结合，通过refine_template提示词逐步优化答案，直至处理完所有相关文本块
    ''','''Refine模式工作原理
    上下文溢出处理‌：若单个文本块过长，会自动分割并继续迭代优化'''])

    print(rs)

def tree_summarize():
    # 初始化TreeSummarize
    summarizer = TreeSummarize(callback_manager=callback_manager)
    # 执行分层摘要
    result = summarizer.get_response(query_str="Refine模式工作原理",text_chunks=['''Refine模式工作原理
    ‌初始回答生成‌：使用第一个检索到的文本块和text_qa_template生成初始回答
    ''','''Refine模式工作原理
    ‌‌迭代优化‌：将初始回答与后续文本块结合，通过refine_template提示词逐步优化答案，直至处理完所有相关文本块
    ''','''Refine模式工作原理
    上下文溢出处理‌：若单个文本块过长，会自动分割并继续迭代优化'''])
    print(result)

if __name__ == '__main__':
    tree_summarize()

def refine():

    simpleSummarize =SimpleSummarize(callback_manager=callback_manager)
    simpleSummarize =Refine(callback_manager=callback_manager)

    rs=simpleSummarize.get_response(query_str="Refine模式工作原理",text_chunks=['''Refine模式工作原理
    ‌初始回答生成‌：使用第一个检索到的文本块和text_qa_template生成初始回答
    ''','''Refine模式工作原理
    ‌‌迭代优化‌：将初始回答与后续文本块结合，通过refine_template提示词逐步优化答案，直至处理完所有相关文本块
    ''','''Refine模式工作原理
    上下文溢出处理‌：若单个文本块过长，会自动分割并继续迭代优化'''])

    print(rs)




